from datetime import datetime
import logging
import os
import json

import numpy as np
import pandas as pd

from .google_dm_api import GoogleDistanceMatrixAPI
from .twogis_dm_api import TWOGISDistanceMatrixAPI
from .constants import OPT_PATH


class DistanceMatrix:
    def __init__(self, config_path='config/distance_matrix_params.json', opt_back_path=OPT_PATH):
        self.logger = logging.getLogger('DistanceMatrix')
        self.logger.debug(f"Initializing DistanceMatrix")
        try:
            with open(os.path.join(opt_back_path, config_path), 'r') as f:
                self.distance_matrix_params = json.load(f)
        except Exception as e:
            self.logger.error(f"Couldn't load config {config_path}, {e}")
            raise FileNotFoundError("Couldn't load config")

        if self.distance_matrix_params['api'] == 'Google':
            self.dm_api = GoogleDistanceMatrixAPI('config/Google_DM_params.json', opt_back_path)
            self.dm_secondary_api = TWOGISDistanceMatrixAPI('config/2GIS_DM_params.json', opt_back_path)
        elif self.distance_matrix_params['api'] == '2GIS':
            self.dm_api = TWOGISDistanceMatrixAPI('config/2GIS_DM_params.json', opt_back_path)
            self.dm_secondary_api = None
        else:
            self.logger.error(f"Unknown DM API name {self.distance_matrix_params['api']}")
            raise KeyError('Unknown DM API name')

        if self.distance_matrix_params['save_path']:
            self.distance_matrix_params['save_path'] = os.path.join(opt_back_path,
                                                                    self.distance_matrix_params['save_path'])
        if self.distance_matrix_params['cache_path']:
            self.distance_matrix_params['cache_path'] = os.path.join(opt_back_path,
                                                                     self.distance_matrix_params['cache_path'])
            self.logger.debug(f"Loading cache: {self.distance_matrix_params['cache_path']}")
            try:
                self.cache = pd.read_csv(self.distance_matrix_params['cache_path'], sep='\t', encoding='utf-8',
                                         dtype={'origin': str, 'destination': str})
                self.cache.set_index(['origin', 'destination'], inplace=True)
            except Exception as e:
                self.logger.error(f"Couldn't load cache file {self.distance_matrix_params['cache_path']}, {e}")
                raise KeyError("Couldn't load cache file")

    @staticmethod
    def coords_to_string(point, precision=3):
        str_point = str(round(np.float32(point[0]), precision)) + '_' + str(round(np.float32(point[1]), precision))
        return str_point

    def update_cache(self, results):
        """Update cache with new origin-destination combinations"""
        self.cache.reset_index(inplace=True)
        self.cache = self.cache[['origin', 'origin_norm', 'destination', 'destination_norm',
                                 'distance', 'duration', 'date', 'source']]
        self.cache = pd.concat([self.cache, results.query('distance != 99.999')], ignore_index=True)
        self.cache.to_csv(self.distance_matrix_params['cache_path'], index=False, sep='\t', encoding='utf-8')
        self.cache.set_index(['origin', 'destination'], inplace=True)

    def find_pair_in_cache(self, origin, destination):
        """Search for origin-destination pair in cache"""
        if not self.distance_matrix_params['cache_path']:
            return None
        origin = self.coords_to_string(origin)
        destination = self.coords_to_string(destination)
        if (origin, destination) in self.cache.index:
            pair = self.cache.loc[(origin, destination)].to_dict('records')[0]
        elif (destination, origin) in self.cache.index:
            pair = self.cache.loc[(destination, origin)].to_dict('records')[0]
        else:
            return None
        pair['origin'] = origin
        pair['destination'] = destination
        return pair

    def build_distance_matrix(self, coordinates, request_params):
        """Build distance matrix of all given coordinates pairs using Google DM API"""
        self.logger.info(f"Building distance matrix")
        results = pd.DataFrame(columns=['origin', 'origin_norm', 'destination', 'destination_norm',
                                        'distance', 'duration', 'date', 'source'])
        cached_rows = []  # Pairs found in cache
        curr_rows = 0
        for i, origin in enumerate(coordinates[:-1]):
            destinations = []
            for destination in coordinates[i + 1:]:
                cached_row = self.find_pair_in_cache(origin, destination)
                if cached_row:
                    cached_rows.append(cached_row)
                else:
                    destinations.append(destination)  # If pair is not found in cache, add it to API calls
                # 25 pairs per request due to Google Maps API limitations
                if len(destinations) == 25:
                    request_result = self.dm_api.get_distance(origin, destinations, request_params)
                    results = pd.concat([results, request_result], ignore_index=True)
                    destinations = []

                curr_rows += 1
                if curr_rows == 1000:
                    if self.distance_matrix_params['cache_path']:
                        self.update_cache(results)
                    curr_rows = 0

            # If destination count can't be divided by 25, we need to handle the remaining pairs
            if len(destinations) != 0:
                request_result = self.dm_api.get_distance(origin, destinations, request_params)
                # print(request_result)
                results = pd.concat([results, request_result], ignore_index=True)
            # break  # Used to limit API calls for debugging

        if self.dm_secondary_api:
            self.logger.info(f"{len(results.query('distance != distance'))} pairs not found by primary DM API")
            secondary__results = []
            for i, row in results.query('distance != distance').iterrows():
                secondary__results.append(self.dm_secondary_api.get_pair_distance(
                    row['origin'], row['destination'], request_params))

            secondary__results_df = pd.DataFrame(secondary__results,
                                          columns=['origin', 'origin_norm', 'destination', 'destination_norm',
                                                   'distance', 'duration', 'date', 'source'])
            results = results.query('distance == distance')
            results = pd.concat([results, secondary__results_df], ignore_index=True)

        self.logger.info(f"Got {len(coordinates)} unique coordinates, {len(results)} new coordinates pairs, "
                         f"{len(results.query('distance != distance'))} pairs not found, "
                         f"{len(results) + len(cached_rows)} total pairs")

        results['origin'] = results['origin'].apply(self.coords_to_string)
        results['destination'] = results['destination'].apply(self.coords_to_string)

        # warn = results.groupby(['destination_norm']).agg({
        #     'duration': lambda x: x.isnull().sum(),
        #     'destination': lambda x: len(x)
        # }).query('duration == destination')
        #
        warnings = []
        # if len(warn):
        #     for dest in warn.index:
        #         warnings.append(f'Адрес {dest} не достижим ни из одной точки. Выбрасываю его.')
        #
        # results = results.loc[~results.destination_norm.isin(warn.index)]

        results['distance'] = results['distance'].fillna(99999)
        results['distance'] = results['distance'] / 1000  # Meters to km
        results['duration'] = results['duration'].fillna(180 * 60)
        results['duration'] = results['duration'] / 60  # Seconds to minutes

        # Combine results from API with cache hits and update cache
        if self.distance_matrix_params['cache_path']:
            self.update_cache(results)
            results = pd.concat([results, pd.DataFrame(cached_rows)], ignore_index=True)

        if self.distance_matrix_params['save_path']:  # Save DM to disk if path is specified
            self.logger.debug(f"Writing distance matrix DataFrame to disk")
            results.to_excel(self.distance_matrix_params['save_path'], index=False, encoding='cp1251')

        return results, warnings


if __name__ == '__main__':
    logging.basicConfig(filename='logs/DistanceMatrix.log',
                        filemode='a',
                        level=logging.DEBUG,
                        format='%(asctime)s,%(msecs)d\t%(name)s\t%(levelname)s\t%(message)s',
                        datefmt='%Y-%m-%d-%H:%M:%S')

    with open('data/input_example_may.json', 'r') as f:
        req_json = json.load(f)
    delivers = pd.DataFrame(req_json['delivers'])
    orders = pd.DataFrame(req_json['orders'])
    meetings = pd.concat([delivers, orders], ignore_index=True, sort=False)
    unique_addresses = list(set(meetings['addr'].values))

    timestamp = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
    dm_builder = DistanceMatrix(config_path='config/distance_matrix_params.json')
    request_params = {'date': datetime.strptime(req_json['date'], '%d.%m.%Y')}
    df = dm_builder.build_distance_matrix(unique_addresses, request_params)
